import numpy as np
from scipy.io import savemat
from scipy.optimize import minimize

# 全局变量声明
f_L = []
g_L = []
k_tmp = 0
AGENT_N_tmp = []
h_tmp = []

class Agent:
    def __init__(self):
        self.joining_step = 0
        self.exiting_step = float('inf')
        self.state_region = np.array([[0, 6], [0, 4]])
        self.malicious_section = []
        self.dynamics = 1

def main_data_generate(T, dt):
    global f_L, g_L, k_tmp, AGENT_N_tmp, h_tmp

    # 初始化智能体列表
    num_agents = 8
    OMAS_L = [Agent() for _ in range(num_agents)]

    # 配置智能体参数
    OMAS_L[0].dynamics = 2
    OMAS_L[1].dynamics = 2
    OMAS_L[4].joining_step = 12
    OMAS_L[5].joining_step = 132
    OMAS_L[6].joining_step = 225
    OMAS_L[3].exiting_step = 3020
    OMAS_L[2].exiting_step = 3320
    OMAS_L[5].exiting_step = 3440
    OMAS_L[5].malicious_section = [[450, 560]]
    OMAS_L[6].malicious_section = [[860, 900], [950, 1100], [2540, 2560]]
    OMAS_L[2].malicious_section = [[2270, 2290], [2700, 2710], [2790, 2900]]

    # 生成AGENT矩阵
    total_steps = int(T/dt) + 1
    AGENT = np.zeros((num_agents, total_steps), dtype=int)

    for k in range(total_steps):
        for i in range(num_agents):
            agent = OMAS_L[i]
            if k >= agent.joining_step and k <= agent.exiting_step:
                AGENT[i, k] = 1  # 默认正常参与

                # 检查恶意时间段
                for section in agent.malicious_section:
                    if len(section) >= 2 and k >= section[0] and k <= section[1]:
                        AGENT[i, k] = 2
                        break

    # 初始化目标函数（需要实现具体函数）
    f_L = [lambda x, k, i=i: function_fL(i, x, k, T, dt) for i in range(num_agents)]
    g_L = [lambda x, k, i=i: function_gL(i, x, k) for i in range(num_agents)]

    # 计算最优解和最优值
    AGENT_opt = [{} for _ in range(num_agents)]
    AGENT_x_opt = [np.full((2, total_steps), np.inf) for _ in range(num_agents)]
    F_opt = np.zeros(total_steps)

    for k in range(total_steps):
        AGENT_tmp = np.where(AGENT[:, k] >= 1)[0].tolist()
        AGENT_N_tmp = np.where(AGENT[:, k] == 1)[0].tolist()

        # print(AGENT_N_tmp)
        
        h_tmp = np.zeros((2, num_agents))
        for idx, i in enumerate(AGENT_N_tmp):
            h_tmp[:, i] = function_hL(idx+1, len(AGENT_tmp), k, dt)
        
        global k_tmp
        k_tmp = k
        
        # 构建初始猜测值
        x_ini = []
        for i in AGENT_N_tmp:
            x_ini_tmp = np.array([1, 1]) + h_tmp[:, i]
            x_ini.extend(x_ini_tmp)
        
        # print(x_ini)
        # 优化求解（需要实现约束函数）
        result = minimize(function_Pglobal, x_ini, constraints=({'type': 'ineq', 'fun': lambda x: function_cglobal(x)[0]},{'type': 'ineq', 'fun': lambda x: function_cglobal(x)[1]}))
        xopt = result.x
        F_opt[k] = result.fun
        
        # 分配最优解
        c = 0
        for i in AGENT_N_tmp:
            AGENT_opt[i][k] = xopt[c:c+2]
            AGENT_x_opt[i][:, k] = xopt[c:c+2]
            c += 2

    # 初始化状态和速度
    AGENT_X = [[np.array([]) for _ in range(total_steps)] for _ in range(num_agents)]
    AGENT_V = [[np.array([]) for _ in range(total_steps)] for _ in range(num_agents)]
    
    for i in range(num_agents):
        join_step = OMAS_L[i].joining_step
        if join_step < total_steps:

            AGENT_X[i][join_step] = AGENT_x_opt[i][:, join_step]
            AGENT_V[i][join_step] = np.array([0.0, 0.0])

    # 生成通信拓扑
    AGENT_A = [np.zeros((num_agents, num_agents)) for _ in range(total_steps)]
    
    for k in range(total_steps):
        AGENT_N_tmp = np.where(AGENT[:, k] == 1)[0].tolist()
        AGENT_M_tmp = np.where(AGENT[:, k] == 2)[0].tolist()

        # print(AGENT_N_tmp,AGENT_M_tmp)
        
        # 创建连接模式
        L_tmp = AGENT_N_tmp + [AGENT_N_tmp[0]]
        L_ttmp = AGENT_N_tmp + AGENT_N_tmp[:2]
        L_tttmp = AGENT_N_tmp + AGENT_N_tmp[:3]
        
        # 正常节点连接
        for i in range(len(AGENT_N_tmp)):
            agent_idx = AGENT_N_tmp[i]
            AGENT_A[k][agent_idx, agent_idx] = 1
            
            # 环形连接
            if i < len(L_tmp)-1:
                AGENT_A[k][L_tmp[i], L_tmp[i+1]] = 1
            
            # 其他连接模式
            # ...（根据实际需求实现）

        # 恶意节点连接
        for i in AGENT_M_tmp:
            AGENT_A[k][i, i] = 1
            for j in AGENT_N_tmp:
                AGENT_A[k][i, j] = 1
                AGENT_A[k][j, i] = 1

    # 保存数据
    data = {
        'OMAS_L': OMAS_L,
        'AGENT': AGENT,
        'AGENT_X': AGENT_X,
        'AGENT_V': AGENT_V,
        'AGENT_opt': AGENT_opt,
        'AGENT_x_opt': AGENT_x_opt,
        'F_opt': F_opt,
        'AGENT_A': AGENT_A
    }
    # print(AGENT_X,AGENT_V)
    
    savemat('data_OMAS_1.mat', data)

# 需要实现的辅助函数占位符
# def function_fL(i, x, k):
#     """本地目标函数f实现"""
#     return 0.0

def function_fL(i, x, k, T, dt):
    """
    本地目标函数f实现
    
    参数:
    i : int - 智能体索引号 (1-based MATLAB索引保持原样)
    x : numpy.ndarray - 智能体状态 (2,)
    k : int - 当前时刻 (时间步索引)
    
    返回:
    float - 目标函数值
    """
    
    # 计算基础参考点
    xc_L = np.array([
        3 + 3 * np.sin(-(np.pi/2) * (4/T) * k * dt),
        2 + 2 * np.sin(-np.pi * (4/T) * k * dt)
    ])
    
    # 添加个体偏移量
    offset = np.array([
        (-0.015 + 0.03 * (i-1)/7) * np.sin(np.pi * k * dt / 40),
        (-0.02 + 0.04 * (i-1)/7) * np.sin(np.pi * k * dt / 20)
    ])
    
    xc_L += offset
    
    # 计算目标函数值
    return 0.5 * np.sum((x - xc_L)**2)

def function_gL(i, x, k):
    """本地目标函数g实现"""
    g=0.05*np.linalg.norm(x,ord=2)
    return g

def function_hL(i, num_AGENT, k, dt):
    """
    编队偏移向量计算函数
    
    参数:
    i : int - 智能体索引
    num_AGENT : int - 当前参与的智能体总数
    k : int - 当前时间步索引
    dt : float - 时间间隔
    
    返回:
    numpy.ndarray - 二维偏移向量 [2,]
    """
    angle = 2 * np.pi * i / num_AGENT + np.pi * k * dt / 80
    h_L = 0.7 * np.array([np.cos(angle), np.sin(angle)])
    return h_L

# def function_Pglobal(x):
#     """全局目标函数实现"""
#     return 0.0

def function_Pglobal(x):
    """
    全局目标函数实现
    
    参数:
    x : numpy.ndarray - 优化变量向量
    
    返回:
    float - 目标函数值
    """
    global f_L, g_L, AGENT_N_tmp, k_tmp
    
    P = 0.0
    c = 0
    
    for i in AGENT_N_tmp:
        # 提取当前智能体的状态变量
        x_i = x[c:c+2]
        
        # 计算f和g的函数值
        f_val = f_L[i](x_i, k_tmp)
        g_val = g_L[i](x_i, k_tmp)

        # print(f_val,g_val,"!!!!!!!")
        
        # 累加到目标函数
        P += f_val + g_val
        
        # 更新变量索引
        c += 2
    
    return P

# def function_cglobal(x):
#     """全局约束函数实现"""
#     return np.array([0.0])

def function_cglobal(x):
    """
    全局约束函数实现
    
    参数:
    x : numpy.ndarray - 优化变量向量
    
    返回:
    tuple - (不等式约束列表, 等式约束列表)
    """
    global AGENT_N_tmp, h_tmp
    
    c = []    # 不等式约束
    ceq = []  # 等式约束
    X = np.array([6, 4])  # 状态约束上限
    
    ci = 0
    for i in AGENT_N_tmp:
        # 提取当前智能体的状态变量
        eta_i_tmp = x[ci:ci+2] - h_tmp[:, i]
        
        cj = 0
        for j in AGENT_N_tmp:
            # 提取邻居智能体的状态变量
            eta_j_tmp = x[cj:cj+2] - h_tmp[:, j]
            
            # 添加等式约束
            ceq.extend((eta_i_tmp - eta_j_tmp).tolist())
            
            # 添加不等式约束
            c.extend((-eta_i_tmp).tolist())            # -eta_i <= 0
            c.extend((eta_i_tmp - X).tolist())         # eta_i <= X
            
            cj += 2  # 每个智能体有2个状态变量
        
        ci += 2  # 移动到下一个智能体的变量
    
    return np.array(c), np.array(ceq)

if __name__ == "__main__":
    T = 40
    dt = 0.01
    main_data_generate(T, dt)
